# train.py - 训练函数模块
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from tqdm import tqdm
import numpy as np
from utils import metrics_of_pv


def train(data, model, criterion, optm, device):
    """训练函数"""
    model.train()
    running_loss = 0.0
    for x, y in tqdm(data):
        model.zero_grad()
        x, y = x.float().to(device), y.float().to(device)
        optm.zero_grad()
        y_pre = model(x)
        loss = criterion(y_pre, y)
        loss.backward()
        optm.step()
        running_loss += loss.item() * x.size(0)
    epoch_loss = running_loss / len(data.dataset)
    return epoch_loss


def evaluate(data, model, criterion, device, scalar=None):
    """
    评估函数 - 用于评估模型在验证集或测试集上的性能
    
    Args:
        data: 数据加载器，提供验证或测试数据批次
        model: 要评估的神经网络模型
        criterion: 损失函数，用于计算预测值与真实值之间的误差
        device: 计算设备（CPU或GPU）
        scalar: 可选的数据标准化器，用于将数据还原到原始尺度
        
    Returns:
        tuple: 包含四个元素的元组
            - epoch_loss (float): 整个数据集的平均损失
            - metrics_ (list): 评估指标列表，包含MAE、MSE、R2和MBE
            - all_labels (numpy.ndarray): 所有真实标签
            - all_preds (numpy.ndarray): 所有预测结果
            
    功能说明:
        1. 将模型设置为评估模式（禁用dropout和batch normalization的随机性）
        2. 遍历数据集，计算每个批次的损失和预测结果
        3. 收集所有预测结果和真实标签
        4. 如果提供了标准化器，将预测结果和真实标签还原到原始尺度
        5. 计算各种评估指标（MAE、MSE、R2、MBE）
        6. 返回平均损失、评估指标、真实标签和预测结果
    """
    # 将模型设置为评估模式，这会禁用dropout和batch normalization的随机性
    model.eval()
    
    # 初始化验证损失累加器，用于计算整个数据集的平均损失
    val_running_loss = 0.0
    
    # 初始化列表，用于存储所有批次的预测结果和真实标签
    all_preds = []
    all_labels = []
    
    # 遍历数据加载器中的每个批次
    for x, y in tqdm(data):  # tqdm用于显示进度条
        # 清空梯度（虽然在评估模式下不需要，但这是一个好习惯）
        model.zero_grad()
        
        # 使用torch.no_grad()上下文管理器，禁用梯度计算
        # 这可以减少内存使用并加速计算，因为在评估阶段不需要计算梯度
        with torch.no_grad():
            # 将输入数据和标签转换为浮点型，并移动到指定设备（CPU或GPU）
            x, y = x.float().to(device), y.float().to(device)
            
            # 使用模型进行前向传播，得到预测结果
            y_pre = model(x)
            
            # 计算损失
            loss = criterion(y_pre, y)
            
            # 累加损失，乘以批次大小以得到正确的加权平均
            val_running_loss += loss.item() * x.size(0)
            
            # 将当前批次的预测结果和真实标签转换为NumPy数组并存储
            all_preds.append(y_pre.cpu().numpy())
            all_labels.append(y.cpu().numpy())
    
    # 计算整个数据集的平均损失
    epoch_loss = val_running_loss / len(data.dataset)

    # 将所有批次的预测结果和真实标签合并成一个数组
    all_preds = np.concatenate(all_preds, axis=0)
    all_labels = np.concatenate(all_labels, axis=0)

    # 如果提供了标准化器，将预测结果和真实标签还原到原始尺度
    # 这对于计算有意义的评估指标很重要，因为模型通常在标准化后的数据上训练
    if scalar is not None:
        all_preds = scalar.inverse_transform(all_preds)
        all_labels = scalar.inverse_transform(all_labels)

    # 计算评估指标（MAE、MSE、R2、MBE）
    metrics_ = metrics_of_pv(all_preds, all_labels)
    
    # 返回平均损失、评估指标、真实标签和预测结果
    return epoch_loss, metrics_, all_labels, all_preds


def create_optimizer_and_scheduler(model, learning_rate):
    """
    创建优化器和学习率调度器
    
    Args:
        model: 神经网络模型，用于获取需要优化的参数
        learning_rate (float): 初始学习率
        
    Returns:
        tuple: 包含优化器和学习率调度器的元组
            - optimizer: Adam优化器实例
            - scheduler: ReduceLROnPlateau学习率调度器实例
            
    功能说明:
        - 创建Adam优化器，使用给定的学习率
        - 创建基于验证损失的学习率调度器，当验证损失不再下降时自动降低学习率
        - 这种组合有助于模型在训练过程中更好地收敛，避免陷入局部最优
    """
    # 创建Adam优化器，使用模型参数和指定的学习率
    # Adam优化器是一种自适应学习率优化算法，结合了动量法和RMSProp的优点
    optimizer = optim.Adam(model.parameters(), lr=learning_rate)
    
    # 创建ReduceLROnPlateau学习率调度器
    # 当验证损失不再改善时，自动降低学习率以帮助模型收敛
    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
        optimizer,        # 要调整学习率的优化器
        mode="min",       # 监控指标的模式，"min"表示指标越小越好
        factor=0.5,       # 学习率衰减因子，新的学习率 = 当前学习率 * factor
        patience=5        # 容忍指标不改善的轮数，超过此值则降低学习率
    )
    
    # 返回优化器和学习率调度器
    return optimizer, scheduler